use sq_global::{API_STOCK_COMP, API_STOCK_DAILY};

use polars::datatypes::DataType;
use polars::prelude::{
    IntoLazy, ListNameSpaceExtension, RoundMode, SortMultipleOptions, col, lit,
};
use tokio::try_join;
use sq_macro::test_with_ctx;
use sq_polars::{get_col_idx, get_str_col};

#[test_with_ctx]
pub async fn test() {
    run().await.expect("run failed");
}

//计算得到最近7个交易日每日连涨大于1%的股票信息
pub async fn run() -> anyhow::Result<Vec<String>> {
    let (com_df, daily_dfs) = try_join!(
        sq_tushare::load_data(API_STOCK_COMP, Vec::new()),
        sq_tushare::recent_dates(API_STOCK_DAILY, 20)
    )?;

    let daily_df = sq_polars::concat_dfs(daily_dfs)?;

    let daily_stats = daily_df
        .lazy()
        .select([col("ts_code"), col("trade_date"), col("pct_chg")])
        .sort(["trade_date"], SortMultipleOptions::default())
        .group_by([col("ts_code")])
        .agg([
            col("pct_chg").tail(Some(7)),
            col("trade_date").tail(Some(7)),
        ])
        .filter(col("pct_chg").list().len().eq(7i64))
        .filter(col("pct_chg").list().min().gt(lit(1.0_f64)))
        .with_column(
            col("pct_chg")
                .list()
                .eval(col("").round(2, RoundMode::HalfAwayFromZero), true),
        );

    let info = com_df
        .lazy()
        .select([col("com_name"), col("ts_code")])
        .inner_join(daily_stats, "ts_code", "ts_code")
        .with_columns([
            col("pct_chg")
                .list()
                .eval(col("").cast(DataType::String) + lit("%"), true)
                .list()
                .join(lit(","), true)
                .alias("continuous_chg_fmt"),
            col("trade_date")
                .list()
                .join(lit(","), true)
                .alias("trade_dates"),
        ])
        .collect()?;

    tracing::info!("info: {:?}", info);

    let map = get_col_idx!(&info);

    let com_name_idx = map["com_name"];
    let ts_code_idx = map["ts_code"];
    let continuous_chg_fmt_idx = map["continuous_chg_fmt"];
    let trade_dates_idx = map["trade_dates"];

    let mut msg_vec = Vec::new();
    for idx in 0..info.height() {
        let row = info.get_row(idx)?;
        let com_name = get_str_col!(row, com_name_idx);
        let ts_code = get_str_col!(row, ts_code_idx);
        let continuous_chg_fmt = get_str_col!(row, continuous_chg_fmt_idx);
        let trade_dates = get_str_col!(row, trade_dates_idx);

        let msg = format!(
            "上市公司{com_name} （股票代码{ts_code}）最近7个交易日{trade_dates}的涨幅为{continuous_chg_fmt}。"
        );
        msg_vec.push(msg);
    }
    
    Ok(msg_vec)
}
